#!/bin/bash

# Download Models Script
# This script downloads YOLO models and converts them to OpenVINO format

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
MODELS_DIR="${1:-${PROJECT_ROOT}/release/goyolopenvino-v1.0.0/models}"
MODEL_TYPE="${2:-yolov8n}"  # yolov8n, yolov8s, yolov11n, etc.

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

echo -e "${GREEN}=== YOLO Model Download and Conversion ===${NC}"
echo "Model Type: $MODEL_TYPE"
echo "Models Directory: $MODELS_DIR"
echo ""

# Create models directory
mkdir -p "$MODELS_DIR"
cd "$MODELS_DIR"

# Determine model URL based on type
case "$MODEL_TYPE" in
    yolov8n)
        MODEL_URL="https://github.com/ultralytics/yolov8/releases/download/v8.0.0/yolov8n.pt"
        MODEL_FILE="yolov8n.pt"
        ;;
    yolov8s)
        MODEL_URL="https://github.com/ultralytics/yolov8/releases/download/v8.0.0/yolov8s.pt"
        MODEL_FILE="yolov8s.pt"
        ;;
    yolov8m)
        MODEL_URL="https://github.com/ultralytics/yolov8/releases/download/v8.0.0/yolov8m.pt"
        MODEL_FILE="yolov8m.pt"
        ;;
    yolov8l)
        MODEL_URL="https://github.com/ultralytics/yolov8/releases/download/v8.0.0/yolov8l.pt"
        MODEL_FILE="yolov8l.pt"
        ;;
    yolov11n)
        MODEL_URL="https://github.com/ultralytics/yolov11/releases/download/v11.0.0/yolov11n.pt"
        MODEL_FILE="yolov11n.pt"
        ;;
    yolov11s)
        MODEL_URL="https://github.com/ultralytics/yolov11/releases/download/v11.0.0/yolov11s.pt"
        MODEL_FILE="yolov11s.pt"
        ;;
    yolov11m)
        MODEL_URL="https://github.com/ultralytics/yolov11/releases/download/v11.0.0/yolov11m.pt"
        MODEL_FILE="yolov11m.pt"
        ;;
    *)
        echo -e "${RED}Error: Unknown model type: $MODEL_TYPE${NC}"
        echo "Supported types: yolov8n, yolov8s, yolov8m, yolov8l, yolov11n, yolov11s, yolov11m"
        exit 1
        ;;
esac

# Download model
echo -e "${YELLOW}Downloading $MODEL_TYPE model...${NC}"
if [ -f "$MODEL_FILE" ]; then
    echo "✓ Model already exists: $MODEL_FILE"
else
    echo "Downloading from: $MODEL_URL"
    wget -q --show-progress "$MODEL_URL" -O "$MODEL_FILE"
    echo "✓ Downloaded: $MODEL_FILE"
fi

# Check if OpenVINO Model Optimizer is available
echo ""
echo -e "${YELLOW}Checking for OpenVINO Model Optimizer...${NC}"

if ! command -v mo &> /dev/null; then
    echo -e "${YELLOW}OpenVINO Model Optimizer not found in PATH${NC}"
    echo ""
    echo "To convert models to OpenVINO format, you need to:"
    echo "1. Install OpenVINO: https://docs.openvino.ai/latest/get_started.html"
    echo "2. Source setupvars.sh:"
    echo "   source /opt/intel/openvino_genai_ubuntu24_2025.3.0.0_x86_64/setupvars.sh"
    echo "3. Run conversion:"
    echo "   mo --input_model $MODEL_FILE --output_dir ."
    echo ""
    echo "Or use the conversion script:"
    echo "   ./scripts/convert-models.sh $MODELS_DIR $MODEL_TYPE"
    exit 0
fi

# Convert to OpenVINO format
echo -e "${YELLOW}Converting to OpenVINO format...${NC}"
MODEL_NAME="${MODEL_FILE%.pt}"

if [ -f "${MODEL_NAME}.xml" ] && [ -f "${MODEL_NAME}.bin" ]; then
    echo "✓ OpenVINO model already exists"
else
    echo "Running: mo --input_model $MODEL_FILE --output_dir ."
    mo --input_model "$MODEL_FILE" --output_dir .
    echo "✓ Converted to OpenVINO format"
fi

# Verify conversion
echo ""
echo -e "${YELLOW}Verifying conversion...${NC}"
if [ -f "${MODEL_NAME}.xml" ] && [ -f "${MODEL_NAME}.bin" ]; then
    echo "✓ ${MODEL_NAME}.xml"
    echo "✓ ${MODEL_NAME}.bin"
    echo ""
    echo -e "${GREEN}=== Conversion Complete ===${NC}"
    echo "Model files:"
    ls -lh "${MODEL_NAME}."*
    echo ""
    echo "Update config/config.env with:"
    echo "  MODEL_PATH=./models/${MODEL_NAME}.xml"
else
    echo -e "${RED}✗ Conversion failed${NC}"
    exit 1
fi

